# set up environment ----
library(data.table)
library(ggplot2)
library(cowplot)
library(nnet)
library(parallel)
load("data/analysis_data.RData")
source("R/00-functions.R")
RNGkind("L'Ecuyer-CMRG")
set.seed(1444055996)

# main dataset ----
data_main <- respondent_data[main == 1]
model_main <- multinom(type2 ~ order, data = data_main)
load("results/boots_main-type2.RData")
predicted_probs_main <- do.call(rbind, lapply(boots_main, predict,
  newdata = data.table(order = data_main[, unique(order)]), type = "probs"))
predicted_probs_main <- as.data.table(predicted_probs_main)
predicted_probs_main[, `:=`(
  order = rep(data_main[, unique(order)], 1000),
  b = rep(1:1000, each = 2))]
levels_main <- get_levels(predicted_probs_main, "order", "Treatment")

# full dataset ----
data_full <- respondent_data
model_full <- multinom(type2 ~ order, data = data_full)
load("results/boots_full-type2.RData")
predicted_probs_full <- do.call(rbind, lapply(boots_full, predict,
  newdata = data.table(order = data_full[, unique(order)]), type = "probs"))
predicted_probs_full <- as.data.table(predicted_probs_full)
predicted_probs_full[, `:=`(
  order = rep(data_full[, unique(order)], 1000),
  b = rep(1:1000, each = 2))]
levels_full <- get_levels(predicted_probs_full, "order", "Treatment")

# plots ----
# treatment/main ----
theme_set(theme_bw() +
    theme(
      legend.position = "none",
      strip.text.x = element_blank(),
      strip.background = element_blank()))
x <- dcast(levels_main[covariate == "Treatment", .(var, type, est)],
  type ~ var, value.var = "est")

g1 <- ggplot(levels_main[covariate == "Treatment"], aes(var, est)) +
  facet_grid( ~ type) +
  geom_segment(data = x, aes(
    x = "Preferences Elicited Last",
    y = `Preferences Elicited Last`,
    xend = "Electability Elicited Last",
    yend = `Electability Elicited Last`)) +
  geom_point() +
  geom_errorbar(aes(ymin = q025, ymax = q975), width = 0.1) +
  xlab("") + ylab("") +
  scale_x_discrete(
    breaks = c("Electability Elicited Last", "Preferences Elicited Last"),
    labels = c("Electability\nLast", "Preferences\nLast")) +
  scale_y_continuous(breaks = seq(.1, .5, .1),
    labels = c("10%", "20%", "30%", "40%", "50%"))+
  theme(strip.background = element_blank())
x1 <- .03
y1 <- 0
y2 <- .12
n <- 1
h <- (1 - y2) / 1
cairo_pdf("figures/type-predictions-levels-treatment-main-type2.pdf", 6.5, 4)
cowplot::ggdraw() +
  draw_plot(g1, x1, y1 + 0 * h, 1 - x1, h) +
  draw_label("Preferences", .25, 1 - y2 * .9, size = 10) +
  draw_label("Electability", .555, 1 - y2 * .9, size = 10) +
  draw_label("Expected Utility", .85, 1 - y2 * .9, size = 10) +
  draw_label("Predicted Probability (Mean and 95% Interval)", x1/2, .5,
    angle = 90, size = 10) +
  draw_label("Predicted Probabilities of Types by Treatment",
    x1 + (1 - x1) / 2, 1 - y2 * .4)
dev.off()

# treatment/full ----
theme_set(theme_bw() +
    theme(
      legend.position = "none",
      strip.text.x = element_blank(),
      strip.background = element_blank()))
x <- dcast(levels_full[covariate == "Treatment", .(var, type, est)],
  type ~ var, value.var = "est")
g1 <- ggplot(levels_full[covariate == "Treatment"], aes(var, est)) +
  facet_grid( ~ type) +
  geom_segment(data = x, aes(
    x = "Preferences Elicited Last",
    y = `Preferences Elicited Last`,
    xend = "Electability Elicited Last",
    yend = `Electability Elicited Last`)) +
  geom_point() +
  geom_errorbar(aes(ymin = q025, ymax = q975), width = 0.1) +
  xlab("") + ylab("") +
  scale_x_discrete(
    breaks = c("Electability Elicited Last", "Preferences Elicited Last"),
    labels = c("Electability\nLast", "Preferences\nLast")) +
  scale_y_continuous(breaks = seq(.1, .5, .1),
    labels = c("10%", "20%", "30%", "40%", "50%"))+
  theme(strip.background = element_blank())
x1 <- .03
y1 <- 0
y2 <- .12
n <- 1
h <- (1 - y2) / 1
cairo_pdf("figures/type-predictions-levels-treatment-full-type2.pdf", 6.5, 4)
cowplot::ggdraw() +
  draw_plot(g1, x1, y1 + 0 * h, 1 - x1, h) +
  draw_label("Preferences", .25, 1 - y2 * .9, size = 10) +
  draw_label("Electability", .555, 1 - y2 * .9, size = 10) +
  draw_label("Expected Utility", .85, 1 - y2 * .9, size = 10) +
  draw_label("Predicted Probability (Mean and 95% Interval)", x1/2, .5,
    angle = 90, size = 10) +
  draw_label("Predicted Probabilities of Types by Treatment (All Respondents)",
    x1 + (1 - x1) / 2, 1 - y2 * .4)
dev.off()

# save ----
save(predicted_probs_main, file = "results/predicted_probs_main-type2.RData")
save(levels_main, file = "results/levels_main-type2.RData")
save(predicted_probs_full, file = "results/predicted_probs_full-type2.RData")
save(levels_full, file = "results/levels_full-type2.RData")
